"""Set of functions to simulate ADCS."""

# %%
import math
from typing import Tuple

import numpy as np
from astropy.constants import G, M_earth
from pyquaternion import Quaternion

from ballistic import sun_vector
from mission_parameters import (
    AVS_MATRIX,
    J_SAT,
    LVLH_O1_GOAL_B_ACCELERATING,
    SOL_GOAL_B,
    WHEEL_ANGLE,
    WHEEL_INERTIA,
    WHEEL_MAX_ACCELERATION,
    WHEEL_MAX_RATE,
)

# M_cmd = − K_d*ω − K_p*q_e
# equation (22) from article
# "Attitude Determination and Control System Design
# for a 6U CubeSat for Proximity Operations and Rendezvous"
# equation (4.2) from article
# Attitude Determination & Control System Design and
# Implementation for a 6U CubeSat Proximity Operations Mission


ANGLE_SIN = math.sin(WHEEL_ANGLE)
ANGLE_COS = math.cos(WHEEL_ANGLE)
# same as math.asin(5/(4*3.2))

WHEELS_MAT = np.array(
    [
        [ANGLE_COS, 0, -ANGLE_COS, 0],
        [0, ANGLE_COS, 0, -ANGLE_COS],
        [ANGLE_SIN, ANGLE_SIN, ANGLE_SIN, ANGLE_SIN],
    ]
)

# Гравитационная постоянная Земли
MU = G.value * M_earth.value


# %%
def angular_velocity_measure(omega_body: np.ndarray) -> np.ndarray:
    """
    Вычисляет угловую скорость спутника в координатах датчика угловой скорости

    Parameters
    ----------
        omega_body: 1x3 np.ndarray координаты угловой скорости спутника в собственной
        системе координат

    Returns
    -------
        omega_sensor: 1x3 np.ndarray - угловая скорость спутника в координатах датчика
        угловой скоростей
    """
    inv_sensor_matrix = np.linalg.inv(AVS_MATRIX)
    omega_sensor = np.array(inv_sensor_matrix.dot(omega_body))[0]
    return omega_sensor


# %%
def quaternion_to_euler_angle(
    q_0: np.ndarray, q_1: np.ndarray, q_2: np.ndarray, q_3: np.ndarray
) -> Tuple[np.ndarray]:
    """Convert Quaternion to Euler Angle

    Parameters
    ----------
        q0 : np.ndarray
        The value vector of the zero component of the quaternion
        q1 : np.ndarray
        The value vector of the first component of the quaternion
        q2 : np.ndarray
        The value vector of the second component of the quaternion
        q3 : np.ndarray
        The value vector of the third component of the quaternion

    Returns
    -------
    res : Tuple[np.ndarray]
        Tuple of Euler angle value vectors
    """
    # https://stackoverflow.com/questions/56207448/efficient-quaternions-to-euler-transformation

    t0 = 2 * (q_0 * q_1 + q_2 * q_3)
    t1 = 1 - 2 * (q_1 * q_1 + q_2 * q_2)
    Euler_angle_0 = np.arctan2(t0, t1)

    t2 = 2 * (q_0 * q_2 - q_3 * q_1)
    t2[t2 > 1] = 1
    t2[t2 < -1] = -1
    Euler_angle_1 = np.arcsin(t2)

    t3 = 2 * (q_0 * q_3 + q_1 * q_2)
    t4 = 1 - 2 * (q_2 * q_2 + q_3 * q_3)
    Euler_angle_2 = np.arctan2(t3, t4)

    return Euler_angle_0, Euler_angle_1, Euler_angle_2


# %%
def gravity_torque(r_earth_sat: np.ndarray, q_bi: np.ndarray) -> np.ndarray:
    """
    Вычисляет гравитационный момент, действующий на КА
    https://stepik.org/lesson/778987/step/3?unit=781512
    https://stepik.org/lesson/804951/step/6?unit=808078

    Parameters:
    ----------
        r_earth_sat: вектор 1x3 радиус вектор КА
        q_bi: вектор 1x4 кватернион ориентации КА из инерциальной в связанную с.к.

    Returns:
    ----------
        Lgr: вектор гравитационного момента
    """

    j_i = -unit_vector(r_earth_sat)
    w = q_bi[0]
    x = q_bi[1]
    y = q_bi[2]
    z = q_bi[3]
    R = np.matrix(
        [
            [1 - 2 * y**2 - 2 * z**2, 2 * x * y + 2 * z * w, 2 * x * z - 2 * y * w],
            [2 * x * y - 2 * z * w, 1 - 2 * x**2 - 2 * z**2, 2 * y * z + 2 * x * w],
            [2 * x * z + 2 * y * w, 2 * y * z - 2 * x * w, 1 - 2 * x**2 - 2 * y**2],
        ]
    )

    j_b = np.array(R.dot(j_i))[0]
    r_earth_sat_norm = np.linalg.norm(r_earth_sat)
    Lgr = np.cross((3 * MU / (r_earth_sat_norm**3)) * (j_b), J_SAT.dot(j_b))

    return Lgr


# %%
def unit_vector(orig_v):
    """Normalize vector

    Parameters
    ----------
    v : numpy array
        original vector

    Returns
    -------
    unit_v : numpy array
        normalized vector
    """

    unit_v = orig_v / np.linalg.norm(orig_v)
    return unit_v


# %%
def a_io(r_i, v_i):
    """The rotation matrix from the orbital frame to the inertial frame.
    Fundamentals of Spacecraft Attitude Determination and Control
    2.6.4 Local-Vertical/Local-Horizontal Frame
    equations 2.78a..2.78c

    Parameters
    ----------
    r_i : 1x3 array vector
        satellite position.
    v_i : 1x3 array vector
        satellite velocity.

    Returns
    -------
    res : 3x3 array
        rotation matrix.
    """

    o3i = -unit_vector(r_i)
    o2i = -unit_vector(np.cross(r_i, v_i))
    o1i = np.cross(o2i, o3i)

    # equation 2.79 in book
    # retVal=cat(dims,o1I,o2I,o3I);
    res = np.concatenate(
        (o1i[np.newaxis].T, o2i[np.newaxis].T, o3i[np.newaxis].T), axis=1
    )

    return res


def omega_mat(omega):
    omega_x, omega_y, omega_z = omega
    # # quaternions according to 2.81 space book, vector than angle q1-3, q4
    # # Initial form (A.22a), (A.1), (A.51) as in space book:
    # omega_matrix = [
    #     [0, omega3, -omega2, omega1],
    #     [-omega3, 0, omega1, omega2],
    #     [omega2, -omega1, 0, omega3],
    #     [-omega1, -omega2, -omega3, 0],
    # ]
    omega_matrix = [
        [0, -omega_x, -omega_y, -omega_z],
        [omega_x, 0, omega_z, -omega_y],
        [omega_y, -omega_z, 0, omega_x],
        [omega_z, omega_y, -omega_x, 0],
    ]
    return omega_matrix


# %%
def wheels_acceleration(cmd_torque, wheels_rate):
    """Determine each wheel acceleration from command torque
    in body frame and actual wheel rates.

    Parameters
    ----------
    cmd_torque : 1x3 list like
        command torque vector
    wheels_rate : 1xn vector
        reaction wheels actual rate

    Returns
    -------
    1xn vector
        reaction wheels accelerations
    """
    wheels_acc = []

    required_rw_torques = np.dot(
        np.dot(
            WHEELS_MAT.T,
            np.linalg.inv(np.dot(WHEELS_MAT, WHEELS_MAT.T)),
        ),
        cmd_torque,
    )

    wheels_acc_required = -1 * required_rw_torques / WHEEL_INERTIA

    for wheel_acc_req, wheel_rate in zip(wheels_acc_required, wheels_rate):
        wheel_acc_i = wheel_acc_req

        if abs(wheel_acc_i) > WHEEL_MAX_ACCELERATION:
            wheel_acc_i = math.copysign(WHEEL_MAX_ACCELERATION, wheel_acc_i)

        if (wheel_rate >= WHEEL_MAX_RATE and wheel_acc_i > 0) or (
            wheel_rate <= -WHEEL_MAX_RATE and wheel_acc_i < 0
        ):
            wheel_acc_i = 0

        wheels_acc.append(wheel_acc_i)

    return np.array(wheels_acc)


# %%
def command_quaternion(t, y, sat_state, strategy):
    """Calculate command quaternion based on satellite state.

    Parameters
    ----------
    t : float
        time, second.
    y : solution vector of arrays for each parameter
        [[internal_energy],
        [m_sum],
        ...
        ]
    sat_state : string
        operational mode name, ex. "eclipse", "heat_accumulation", etc.

    Returns
    -------
    Quaternion object
        command quaternion
    """

    if strategy == "increase periapsis":
        LVLH_O1_GOAL_B = LVLH_O1_GOAL_B_ACCELERATING
    else:
        LVLH_O1_GOAL_B = -1 * np.array(LVLH_O1_GOAL_B_ACCELERATING)

    if sat_state == "heat_accumulation":
        # r_earth_sun = sun_vector(t)
        # e_sun_eci = r_earth_sun / np.linalg.norm(r_earth_sun)
        e_sun_eci = sun_vector(t)

        phi = math.acos(np.dot(e_sun_eci, SOL_GOAL_B))
        quat_cmd_v_e = unit_vector(np.cross(e_sun_eci, SOL_GOAL_B))
        quat_cmd_v = quat_cmd_v_e * math.sin(phi / 2)
        quat_cmd_0 = math.cos(phi / 2)
        quat_cmd = Quaternion(scalar=quat_cmd_0, vector=quat_cmd_v)
    elif sat_state in (
        "orbital_orientation",
        "thrusting",
        "waiting_thrusting_allowed_zone",
    ):
        r_vector = y[2:5]
        v_vector = y[5:8]
        o1_eci = a_io(r_vector, v_vector)[:, 0]

        phi = math.acos(np.dot(o1_eci, LVLH_O1_GOAL_B))
        quat_cmd_v_e = unit_vector(np.cross(o1_eci, LVLH_O1_GOAL_B))
        quat_cmd_v = quat_cmd_v_e * math.sin(phi / 2)
        quat_cmd_0 = math.cos(phi / 2)
        quat_cmd = Quaternion(scalar=quat_cmd_0, vector=quat_cmd_v).unit
    elif sat_state in ("eclipse", "out_of_fuel"):
        # q_bi = y[11:15]
        # quat_cmd = Quaternion(q_bi)  # leave quaternion unchanged
        quat_cmd = Quaternion([1, 0, 0, 0])
        # K_P = 0, quat_cmd has no influence
    else:
        raise Exception("Unknown satellite state")
    return quat_cmd


def error_quaternion(q_bi, q_cmd):
    """Calculate error quaternion from actual and command quaternion

    Parameters
    ----------
    q_bi : Quaternion object
        Rotation quaternion from inertial to body reference frame
    q_cmd : Quaternion object
        Command quaternion

    Returns
    -------
    Quaternion object
        Error quaternion
    """
    v0 = np.array(q_bi.elements)
    v1 = np.array(q_cmd.elements)

    # if quat_err.scalar < 0:
    #     quat_err = -quat_err
    # https://en.wikipedia.org/wiki/Slerp
    # https://stackoverflow.com/questions/21513637/dot-product-of-two-quaternion-rotations
    # eq. 2.8 in PhD Attitude determination ... for 6U
    # note difference between \bigotimes and \bigodot multiplication
    # eq. 2.83 in book Fundamentals of spacecraft attitude determination ...
    # q1 \bigotimes q2 == q2 \bigodot q1

    if np.sum(v0 * v1) < 0:
        q_bi = -q_bi
    # quat_err = q_bi * q_cmd.inverse
    # BUG (possibly) difference with PhD article
    quat_err = q_cmd.inverse * q_bi
    return quat_err


# %%
# Space book: equation 3.147, 3.81
